import pandas as pd
from sklearn.ensemble import RandomForestRegressor
import os

def predict_workload(year, week_number):
    # Загрузка и подготовка данных
    studies_df = pd.read_csv('backend/ml/historical_studies_data.csv')

    # Преобразование данных в числовые признаки
    studies_df['is_with_contrast'] = studies_df['is_with_contrast'].astype(int)
    studies_df['is_multiple_zones'] = studies_df['is_multiple_zones'].astype(int)

    # Создание новых признаков (например, сезонность)
    studies_df['date'] = pd.to_datetime(studies_df['year'].astype(str) + studies_df['week_number'].astype(str) + '1', format='%Y%W%w')
    studies_df['month'] = studies_df['date'].dt.month
    studies_df['day_of_week'] = studies_df['date'].dt.dayofweek

    # Разделение данных на признаки и целевые переменные
    features = studies_df[['year', 'week_number', 'modality', 'is_with_contrast', 'is_multiple_zones', 'month', 'day_of_week']]
    target_quantity = studies_df['quantity']
    target_ue = studies_df['ue_per_study'] * studies_df['quantity']

    # Преобразование категориальных признаков в dummy-переменные
    features = pd.get_dummies(features)

    # Создание и обучение модели прогнозирования
    model_quantity = RandomForestRegressor()
    model_quantity.fit(features, target_quantity)

    model_ue = RandomForestRegressor()
    model_ue.fit(features, target_ue)

    # Прогнозирование на основе входных данных
    input_data = {
        'year': [year] * 13,
        'week_number': [week_number] * 13,
        'modality': ['Денситометр', 'КТ', 'КТ с КУ 1 зона', 'КТ с КУ 2 и более зон', 'ММГ', 'МРТ', 'МРТ с КУ 1 зона', 'МРТ с КУ 2 и более зон', 'РГ', 'Флюорограф'],
        'is_with_contrast': [0, 0, 1, 1, 0, 0, 1, 1, 0, 0],
        'is_multiple_zones': [0, 0, 0, 1, 0, 0, 0, 1, 0, 0],
        'month': [pd.to_datetime(f'{year}-{week_number}-1', format='%Y-%W-%w').month] * 13,
        'day_of_week': [pd.to_datetime(f'{year}-{week_number}-1', format='%Y-%W-%w').dayofweek] * 13
    }
    input_df = pd.DataFrame(input_data)
    input_df = pd.get_dummies(input_df)
    input_df = input_df.reindex(columns=features.columns, fill_value=0)

    predicted_quantity = model_quantity.predict(input_df)
    predicted_ue = model_ue.predict(input_df)

    predicted_workload = [
        {
            'modality': modality,
            'predicted_quantity': int(quantity),
            'predicted_ue': float(ue)
        }
        for modality, quantity, ue in zip(input_data['modality'], predicted_quantity, predicted_ue)
    ]

    return predicted_workload

